from langchain_openai import ChatOpenAI
from jwt_token import generate_token
from langchain_community.llms import Baichuan
from langchain_community.llms import QianfanLLMEndpoint
from langchain_community.llms import Ollama
import os

class MyLLM:
    def __init__(self, model_name, temprature=0.7):
        self.model_name = model_name
        self.temprature = temprature

    def get_llm(self):
        if self.model_name == "gpt-4":
            llm = ChatOpenAI(
                model_name="gpt-4",
                openai_api_key=generate_token(os.environ["OPENAI_API_KEY"], 3600),
                streaming=False,
                temperature=self.temprature
            )
            return llm
        elif self.model_name == "glm-4":
            llm = ChatOpenAI(
                model_name="glm-4",
                openai_api_base="https://open.bigmodel.cn/api/paas/v4",
                openai_api_key=generate_token(os.environ["ZHIPU_API_KEY"], 3600),
                streaming=False,
                temperature=self.temprature
            )
            return llm
        elif self.model_name == "baichuan":
            llm = Baichuan(baichuan_api_key=os.environ["BAICHUAN_API_KEY"])
            return llm
        elif self.model_name == "qianfan":
            llm = QianfanLLMEndpoint(streaming=True,model="ERNIE-Bot-4")
            return llm
        elif self.model_name == "mistral":
            llm = Ollama(model="mistral")
            return llm
        elif self.model_name == "gemma":
            llm = Ollama(model="gemma:2b")
            return llm
        

# llm = MyLLM(model_name="qianfan").get_llm()
# print(llm.model)